#!/usr/bin/env python3
"""
检查EchoNet-Dynamic数据集的预处理状态
判断是否还需要继续预处理
"""

import json
import sys
from pathlib import Path
from typing import Dict, List, Set

def check_echonet_status():
    """检查EchoNet-Dynamic数据集的预处理状态"""
    
    # 源数据目录
    source_dir = Path("D:/Data/EchoNet-Dynamic/Videos")
    
    # 已处理数据目录
    processed_dir = Path("data/processed/segmentation/echonet_dynamic")
    
    # 预处理结果文件
    results_file = Path("data/processed/segmentation/preprocessing_results.json")
    
    print("="*60, flush=True)
    print("EchoNet-Dynamic数据集预处理状态检查", flush=True)
    print("="*60, flush=True)
    print(flush=True)
    
    # 1. 检查源数据
    if source_dir.exists():
        video_files = list(source_dir.glob("*.avi"))
        print(f"源数据目录: {source_dir}")
        print(f"  视频文件总数: {len(video_files)}")
    else:
        print(f"⚠️  源数据目录不存在: {source_dir}")
        video_files = []
    
    # 2. 检查已处理的数据
    processed_videos = set()
    if processed_dir.exists():
        processed_dirs = [d for d in processed_dir.iterdir() if d.is_dir()]
        processed_videos = {d.name for d in processed_dirs}
        print(f"\n已处理数据目录: {processed_dir}")
        print(f"  已处理视频数: {len(processed_videos)}")
        
        # 检查文件完整性
        complete_count = 0
        incomplete_count = 0
        for video_id in processed_videos:
            video_dir = processed_dir / video_id
            has_frames = any(video_dir.glob("*.png"))
            has_metadata = (video_dir / "metadata.json").exists()
            if has_frames and has_metadata:
                complete_count += 1
            else:
                incomplete_count += 1
        
        print(f"  完整处理: {complete_count} 个")
        print(f"  不完整: {incomplete_count} 个")
    else:
        print(f"\n⚠️  已处理数据目录不存在: {processed_dir}")
    
    # 3. 检查预处理结果文件
    if results_file.exists():
        with open(results_file, 'r', encoding='utf-8') as f:
            results = json.load(f)
        
        echonet_result = results.get('echonet_dynamic', {})
        total_samples = echonet_result.get('total_samples', 0)
        valid_samples = echonet_result.get('valid_samples', 0)
        samples = echonet_result.get('samples', [])
        
        print(f"\n预处理结果文件: {results_file}")
        print(f"  记录的总样本数: {total_samples}")
        print(f"  有效样本数: {valid_samples}")
        print(f"  实际样本数: {len(samples)}")
        
        # 从结果中提取已处理的视频ID
        result_video_ids = {s.get('id', '') for s in samples}
        print(f"  结果中的视频ID数: {len(result_video_ids)}")
    else:
        print(f"\n⚠️  预处理结果文件不存在: {results_file}")
        result_video_ids = set()
    
    # 4. 检查数据划分文件
    split_files = {
        'train': Path("data/processed/segmentation/train_split.json"),
        'val': Path("data/processed/segmentation/val_split.json"),
        'test': Path("data/processed/segmentation/test_split.json")
    }
    
    split_video_ids = set()
    split_samples_by_type = {'echonet_dynamic': 0, 'cardiacnet': 0, 'other': 0}
    
    print(f"\n数据划分文件检查:")
    for split_name, split_file in split_files.items():
        if split_file.exists():
            with open(split_file, 'r', encoding='utf-8') as f:
                split_data = json.load(f)
            samples = split_data.get('samples', [])
            echonet_samples = [s for s in samples if s.get('type') == 'echonet_dynamic']
            split_video_ids.update(s.get('id', '') for s in echonet_samples)
            split_samples_by_type['echonet_dynamic'] += len(echonet_samples)
            split_samples_by_type['cardiacnet'] += sum(1 for s in samples if s.get('type') == 'cardiacnet')
            split_samples_by_type['other'] += sum(1 for s in samples if s.get('type') not in ['echonet_dynamic', 'cardiacnet'])
            print(f"  {split_name}: {len(samples)} 个样本 (EchoNet: {len(echonet_samples)})")
        else:
            print(f"  {split_name}: 文件不存在")
    
    total_split_samples = sum(split_samples_by_type.values())
    print(f"\n数据划分统计:")
    print(f"  总样本数: {total_split_samples}")
    print(f"  EchoNet-Dynamic: {split_samples_by_type['echonet_dynamic']} 个")
    print(f"  CardiacNet: {split_samples_by_type['cardiacnet']} 个")
    print(f"  其他: {split_samples_by_type['other']} 个")
    print(f"  数据划分中的EchoNet视频ID数: {len(split_video_ids)}")
    
    # 5. 分析需要处理的视频
    if video_files:
        source_video_ids = {f.stem for f in video_files}
        print(f"\n源数据视频ID数: {len(source_video_ids)}")
        
        # 找出未处理的视频（未在已处理目录中，也未在结果文件中）
        unprocessed = source_video_ids - processed_videos - result_video_ids
        print(f"\n未处理的视频数: {len(unprocessed)}")
        
        # 找出已处理但未在数据划分中的视频
        processed_not_in_splits = processed_videos - split_video_ids
        print(f"已处理但未在数据划分中的视频数: {len(processed_not_in_splits)}")
        
        if unprocessed:
            print(f"  前10个未处理的视频ID:")
            for i, vid_id in enumerate(sorted(unprocessed)[:10], 1):
                print(f"    {i}. {vid_id}")
            if len(unprocessed) > 10:
                print(f"    ... 还有 {len(unprocessed) - 10} 个")
        
        # 计算处理进度
        processed_count = len(processed_videos | result_video_ids)
        progress = (processed_count / len(source_video_ids) * 100) if source_video_ids else 0
        print(f"\n处理进度: {processed_count}/{len(source_video_ids)} ({progress:.1f}%)")
        
        # 判断是否需要继续预处理
        print("\n" + "="*60)
        needs_reprocessing = False
        reasons = []
        
        if len(unprocessed) > 0:
            needs_reprocessing = True
            reasons.append(f"还有 {len(unprocessed)} 个视频未处理")
        
        if len(processed_not_in_splits) > 0 and len(processed_videos) > 50:
            needs_reprocessing = True
            reasons.append(f"有 {len(processed_not_in_splits)} 个已处理视频未包含在数据划分中")
            reasons.append("建议: 使用 --full_mode 重新运行预处理以更新数据划分")
        
        if len(result_video_ids) < len(processed_videos) and len(processed_videos) > 50:
            needs_reprocessing = True
            reasons.append(f"预处理结果文件只记录了 {len(result_video_ids)} 个样本，但实际已处理 {len(processed_videos)} 个")
            reasons.append("建议: 使用 --full_mode 重新运行预处理以更新结果文件")
        
        if needs_reprocessing:
            print("⚠️  结论: EchoNet-Dynamic数据集需要重新运行预处理")
            for reason in reasons:
                print(f"   - {reason}")
        else:
            print("✅ 结论: EchoNet-Dynamic数据集已全部处理完成")
            if len(processed_videos) == len(source_video_ids):
                print("   所有视频都已处理完成")
            if len(split_video_ids) > 0:
                print(f"   数据划分中包含 {len(split_video_ids)} 个EchoNet样本")
        print("="*60)
    else:
        print("\n⚠️  无法检查源数据，无法判断是否需要继续预处理")

if __name__ == "__main__":
    check_echonet_status()

